import os
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import pandas as pd
from pathlib import Path
from Crypto.Util.Padding import pad


# Constants matching your C code
KEY_SIZE = 32  # AES-256
IV_SIZE = 16   # AES block size
MAX_LINE_LENGTH = 256



def parse_excel_file():
    # 读取Excel文件
    file_path = 'config.xlsx'

    # 获取所有工作表名称
    sheets = pd.ExcelFile(file_path).sheet_names
    print(f"工作表列表: {sheets}")

    # 解析Sheet1
    sheet1_data = pd.read_excel(file_path, sheet_name='Sheet1')
    print("\nSheet1内容:")
    print(sheet1_data)

    # 解码ACCESS_KEY
    print("\n解码后的ACCESS_KEY:")
    # 遍历每一行数据
    # 遍历每一行数据
    for index, row in sheet1_data.iterrows():
        product_id = row['PRODUCT_ID']
        device_name = row['DEVICE_NAME']
        access_key = row['ACCESS_KEY']
        
        # 创建目录结构: 当前目录/PRODUCT_ID/DEVICE_NAME/
        dir_path = Path(f"{product_id}/{device_name}")
        dir_path.mkdir(parents=True, exist_ok=True)  # 自动创建不存在的目录
        
        # 创建config.txt文件路径
        config_file = dir_path / "config.txt"
        output_file = dir_path / "config.enc"
        iv_path = dir_path / "config.iv" 

        # Generate random IV
        iv = get_random_bytes(IV_SIZE)

        # 写入内容到文件
        with open(config_file, 'w') as f:
            f.write(f'PRODUCT_ID="{product_id}"\n')
            f.write(f'DEVICE_NAME="{device_name}"\n')
            f.write(f'ACCESS_KEY="{access_key}"')
        
        print(f"已创建: {config_file}")
        encrypt_file(config_file, output_file, iv)
        # Save IV to separate file (needed for decryption)
        with open(iv_path, 'wb') as iv_file:
            iv_file.write(iv)

        print("所有数据已按结构保存完成")

    # 检查其他工作表是否为空
    # for sheet in sheets[1:]:  # 从第二个工作表开始检查
    #     df = pd.read_excel(file_path, sheet_name=sheet)
    #     if df.empty:
    #         print(f"\n{sheet} 是空工作表")
    #     else:
    #         print(f"\n{sheet} 内容:")
    #         print(df)

def encrypt_file(input_file, output_file, iv):
        
    # In a real application, get this from a secure source
    key = b"my_secret_key_32_bytes_long_1234"  # Must be 32 bytes

    cipher = AES.new(key, AES.MODE_CBC, iv)
    
    with open(input_file, 'rb') as fin, open(output_file, 'wb') as fout:
        while True:
            chunk = fin.read(MAX_LINE_LENGTH)
            if not chunk:
                break
                
            # Pad the last chunk if needed
            # if len(chunk) % AES.block_size != 0:
            #     chunk += b' ' * (AES.block_size - len(chunk) % AES.block_size)
            # pad(chunk, AES.block_size)  
            # encrypted = cipher.encrypt(chunk)
            encrypted = cipher.encrypt(pad(chunk, AES.block_size))  # Auto PKCS#7
            fout.write(encrypted)

def encrypt_file_no_pad(input_file, output_file, iv):
       # In a real application, get this from a secure source
    key = b"my_secret_key_32_bytes_long_1234"  # Must be 32 bytes
    # 检查密钥和IV长度
    if len(key) != 32 or len(iv) != 16:
        raise ValueError("Key must be 32 bytes, IV must be 16 bytes")
    
    cipher = AES.new(key, AES.MODE_CBC, iv)
    with open(input_file, 'rb') as fin, open(output_file, 'wb') as fout:
        plaintext = fin.read()
        if len(plaintext) % AES.block_size != 0:
            raise ValueError("Input file size must be a multiple of 16 bytes")
        fout.write(cipher.encrypt(plaintext))

if __name__ == "__main__":
    parse_excel_file()